import re
import json
from typing import Dict, Any, List
from .medical_record import MedicalRecord


class MedicalRecordExtractor:
    """
    电子病历信息提取器
    从病历文本中提取关键信息
    """
    
    def extract_info(self, record_data: Dict[str, Any]) -> MedicalRecord:
        """
        从病历数据中提取关键信息
        
        Args:
            record_data (Dict[str, Any]): 病历原始数据
            
        Returns:
            MedicalRecord: 提取后的病历对象
        """
        # 提取患者ID
        patient_id = self._extract_patient_id(record_data)
        
        # 提取诊断结果
        diagnosis = self._extract_diagnosis(record_data)
        
        # 提取用药清单
        medications = self._extract_medications(record_data)
        
        # 提取检查报告日期
        report_date = self._extract_report_date(record_data)
        
        return MedicalRecord(
            patient_id=patient_id,
            diagnosis=diagnosis,
            medications=medications,
            report_date=report_date,
            raw_data=json.dumps(record_data)
        )
    
    def _extract_patient_id(self, record_data: Dict[str, Any]) -> str:
        """
        提取患者ID
        
        Args:
            record_data (Dict[str, Any]): 病历数据
            
        Returns:
            str: 患者ID
        """
        # 尝试从不同字段获取患者ID
        patient_id = record_data.get('patient_id') or record_data.get('patientId') or record_data.get('id', '')
        
        # 清理文本
        if isinstance(patient_id, str):
            patient_id = re.sub(r'\s+', ' ', patient_id.strip())
        else:
            patient_id = str(patient_id)
            
        return patient_id
    
    def _extract_diagnosis(self, record_data: Dict[str, Any]) -> str:
        """
        提取诊断结果
        
        Args:
            record_data (Dict[str, Any]): 病历数据
            
        Returns:
            str: 诊断结果
        """
        # 尝试从不同字段获取诊断结果
        diagnosis = record_data.get('diagnosis') or record_data.get('diagnose') or record_data.get('result', '')
        
        # 清理文本
        if isinstance(diagnosis, str):
            diagnosis = re.sub(r'\s+', ' ', diagnosis.strip())
        else:
            diagnosis = str(diagnosis)
            
        return diagnosis
    
    def _extract_medications(self, record_data: Dict[str, Any]) -> List[str]:
        """
        提取用药清单
        
        Args:
            record_data (Dict[str, Any]): 病历数据
            
        Returns:
            List[str]: 用药清单
        """
        # 尝试从不同字段获取用药清单
        medications = record_data.get('medications') or record_data.get('drugs') or record_data.get('medication', [])
        
        # 处理不同格式
        if isinstance(medications, str):
            try:
                # 尝试解析JSON格式
                medications = json.loads(medications)
            except json.JSONDecodeError:
                # 如果不是JSON格式，按逗号分割
                medications = [med.strip() for med in medications.split(',') if med.strip()]
        elif not isinstance(medications, list):
            medications = [str(medications)]
            
        # 清理每个药物名称
        cleaned_medications = []
        for med in medications:
            if isinstance(med, str):
                cleaned_med = re.sub(r'\s+', ' ', med.strip())
                if cleaned_med:  # 只添加非空药物
                    cleaned_medications.append(cleaned_med)
            else:
                cleaned_medications.append(str(med))
                
        return cleaned_medications
    
    def _extract_report_date(self, record_data: Dict[str, Any]) -> str:
        """
        提取检查报告日期
        
        Args:
            record_data (Dict[str, Any]): 病历数据
            
        Returns:
            str: 检查报告日期
        """
        # 尝试从不同字段获取检查报告日期
        report_date = record_data.get('report_date') or record_data.get('date') or record_data.get('reportDate', '')
        
        # 清理文本
        if isinstance(report_date, str):
            report_date = re.sub(r'\s+', ' ', report_date.strip())
        else:
            report_date = str(report_date)
            
        return report_date